import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
import numpy as np
from typing import List, Dict, Optional, Union, Literal

from matplotlib.colors import LinearSegmentedColormap

# 全局样式设置
def set_global_style() -> None:
    """设置全局绘图样式"""
    sns.set_theme(style="whitegrid")
    plt.rcParams['font.sans-serif'] = ['SimHei']  # 解决中文显示
    plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示

# 柱状图绘制
def draw_barplot(
    data: pd.DataFrame,
    x_col: str,
    y_col: str,
    hue_col: Optional[str],
    title: str,
    save_path: str,
    palette: Union[str, List[str]] = "viridis",
    rotation: int = 45,
    figsize: tuple = (14, 8)
) -> None:
    """
    绘制分组柱状图
    
    参数:
    data: 包含绘图数据的DataFrame
    x_col: X轴列名
    y_col: Y轴列名
    hue_col: 分组列名(可选)
    title: 图表标题
    save_path: 图片保存路径
    palette: Seaborn调色板名称或颜色列表
    rotation: X轴标签旋转角度
    figsize: 图像尺寸
    """
    set_global_style()
    plt.figure(figsize=figsize)
    
    ax = sns.barplot(
        data=data,
        x=x_col,
        y=y_col,
        hue=hue_col,
        palette=palette,
        errorbar=None
    )
    
    # plt.title(title, fontsize=15)
    plt.xticks(rotation=rotation)
    plt.axhline(0, color='gray', linestyle='--', alpha=0.7)
    
    if hue_col:
        plt.legend(title=hue_col, title_fontsize=12)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.show()

# 热力图绘制
def draw_heatmap(
    data: pd.DataFrame,
    title: str,
    save_path: str,
    cmap: Optional[Union[str, LinearSegmentedColormap]],
    annot: bool = True,
    fmt: str = ".2f",
    figsize: tuple = (12, 10)
) -> None:
    """
    绘制相关系数热力图
    
    参数:
    data: 数值型DataFrame
    title: 图表标题
    save_path: 图片保存路径
    cmap: 颜色映射
    annot: 是否显示数值
    fmt: 数值格式
    figsize: 图像尺寸
    """
    set_global_style()
    plt.figure(figsize=figsize)
    
    sns.heatmap(
        data=data,
        cmap=cmap,
        annot=annot,
        fmt=fmt,
        linewidths=0.5,
        center=0
    )
    
    # plt.title(title, fontsize=15)
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.show()

# 箱线图绘制
def draw_boxplot(
    data: pd.DataFrame,
    x_col: str,
    y_col: str,
    hue_col: Optional[str],
    title: str,
    save_path: str,
    palette: Union[str, List[str]] = "Set2",
    figsize: tuple = (12, 8)
) -> None:
    """
    绘制分组箱线图
    
    参数:
    data: 包含绘图数据的DataFrame
    x_col: X轴列名
    y_col: Y轴列名
    hue_col: 分组列名(可选)
    title: 图表标题
    save_path: 图片保存路径
    palette: Seaborn调色板名称或颜色列表
    figsize: 图像尺寸
    """
    set_global_style()
    plt.figure(figsize=figsize)
    
    sns.boxplot(
        data=data,
        x=x_col,
        y=y_col,
        hue=hue_col,
        palette=palette,
        showfliers=False  # 不显示异常值
    )
    
    # plt.title(title, fontsize=15)
    plt.xticks(rotation=45)
    plt.grid(True, alpha=0.3)
    
    if hue_col:
        plt.legend(title=hue_col, title_fontsize=12)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.show()

heavy_blue : str = "#5A9AD4"
heavy_purple : str = "#8A8AE4"
heavy_green : str = "#72AF47"
heavy_orange: str = "#FF9655"
heavy_red: str = "#EA6B66"
heavy_gray: str = "#A8A8A8"

light_blue: str = "#B8D8F8"    
light_purple: str = "#D4D4FF"  
light_green: str = "#C8E6A8"   
light_orange: str = "#FFD8A8"  
light_red: str = "#FFB6B3"     
light_gray: str = "#E0E0E0"    

styles_count:int = 5

color_schemes:dict [str, list[str]] = {
    "heavy_1": [heavy_blue, heavy_purple, heavy_green, heavy_orange, heavy_red],
    "light_1": [light_blue, light_purple, light_green, light_red, light_orange]
}

def auto_barplot(
    data: pd.DataFrame,
    x_col: str,
    y_col: str,
    hue_col: Optional[str],
    save_path: str,
    rotation: int = 45,
    figsize: tuple = (14, 8),
    title: str = ""
) -> None:
    """
    自动配色的分组柱状图
    """
    # 根据分组数量自动选择配色
    n_group = data[hue_col].nunique() if hue_col else 1
    palette = color_schemes["light_1"] if n_group <= styles_count else "tab10"
    draw_barplot(
        data=data,
        x_col=x_col,
        y_col=y_col,
        hue_col=hue_col,
        title=title,
        save_path=save_path,
        palette=palette,
        rotation=rotation,
        figsize=figsize
    )

def auto_boxplot(
    data: pd.DataFrame,
    x_col: str,
    y_col: str,
    hue_col: Optional[str],
    save_path: str,
    figsize: tuple = (12, 8),

    title: str = "",
) -> None:
    """
    自动配色的分组箱线图
    """
    n_group = data[hue_col].nunique() if hue_col else 1
    palette = color_schemes["heavy_1"] if n_group <= styles_count else "tab10"
    draw_boxplot(
        data=data,
        x_col=x_col,
        y_col=y_col,
        hue_col=hue_col,
        title=title,
        save_path=save_path,
        palette=palette,
        figsize=figsize,
        
    )

# 饼状图绘制
def draw_pieplot(
    data: pd.Series,
    title: str,
    save_path: str,
    colors: Optional[List[str]] = None,
    autopct: str = '%1.1f%%',
    startangle: int = 90,
    figsize: tuple = (10, 8),
    explode: Optional[List[float]] = None
) -> None:
    """
    绘制饼状图
    
    参数:
    data: 包含绘图数据的Series，index为标签，values为数值
    title: 图表标题
    save_path: 图片保存路径
    colors: 自定义颜色列表
    autopct: 百分比显示格式
    startangle: 起始角度
    figsize: 图像尺寸
    explode: 扇形分离程度列表
    """
    set_global_style()
    plt.figure(figsize=figsize)
    
    result = plt.pie(
        data.values.tolist(),
        labels=data.index.tolist(),
        colors=colors,
        autopct=autopct,
        startangle=startangle,
        explode=explode,
        shadow=False
    )
    wedges, _ = result[:2]
    # _ = result[2] if len(result) > 2 else []
    
    plt.title(title, fontsize=15)
    plt.axis('equal')  # 确保饼图是圆的

    # 添加图例，显示颜色与标签的对应关系
    plt.legend(wedges, data.index.tolist(), title="类别", loc="best", bbox_to_anchor=(1, 0.5))

    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.show()

def auto_pieplot(
    data: pd.Series,
    save_path: str,
    title: str = "",
    autopct: str = '%1.1f%%',
    startangle: int = 90,
    figsize: tuple = (10, 8),
    explode: Optional[List[float]] = None
) -> None:
    """
    自动配色的饼状图
    """
    # 根据数据项数量自动选择配色
    n_items = len(data)
    colors = color_schemes["heavy_1"] if n_items <= styles_count else None
    
    draw_pieplot(
        data=data,
        title=title,
        save_path=save_path,
        colors=colors,
        autopct=autopct,
        startangle=startangle,
        figsize=figsize,
        explode=explode
    )


def draw_distribution_fit(
    data: pd.Series,
    fit_func,
    fit_params: tuple,
    title: str,
    save_path: str,
    bins: int = 30,
    alpha: float = 0.6,
    figsize: tuple = (10, 6),
    label_data: str = '数据分布',
    label_fit: str = '拟合分布'
) -> None:
    """
    绘制数据分布直方图和拟合曲线
    
    参数:
    data: 包含数据的Series
    fit_func: 拟合分布的概率密度函数 (如 stats.norm.pdf)
    fit_params: 拟合参数元组 (如 (mu, sigma))
    title: 图表标题
    save_path: 图片保存路径
    bins: 直方图分箱数
    alpha: 直方图透明度
    figsize: 图像尺寸
    label_data, label_fit: 图例标签
    """
    set_global_style()
    fig, ax = plt.subplots(figsize=figsize)
    
    # 绘制直方图 - 使用 heavy_blue 颜色
    ax.hist(data, bins=bins, density=True, alpha=alpha, color=heavy_blue, label=label_data)
    
    # 绘制拟合曲线 - 使用虚线
    x = np.linspace(data.min(), data.max(), 200)
    y = fit_func(x, *fit_params)
    ax.plot(x, y, 'k--', lw=2, label=label_fit)
    
    # ax.set_title(title)
    ax.legend()
    ax.grid(True, alpha=0.3)
    
    plt.tight_layout()
    plt.savefig(save_path, dpi=300)
    plt.show()


# 蓝紫渐变
custom_blue_purple = LinearSegmentedColormap.from_list(
    "blue_purple", ["#FBF9FA", "#DEDCEA", "#AFB7DB", "#8197C6", "#507AAF"], N=256
)

# 橙红渐变
custom_orange_red = LinearSegmentedColormap.from_list(
    "orange_red", ["#F7F7E9", "#F3E1AF", "#DBBF92", "#D78851", "#BE5C37"], N=256
)

# 绿色系
custom_green = LinearSegmentedColormap.from_list(
    "green", ["#F8FBF6", "#DEEED4", "#BBDCA1", "#86C06C", "#5B9C4B"], N=256
)

# 蓝红渐变

custom_blue_red = LinearSegmentedColormap.from_list(
    "blue_red", ["#053061", "#134b87", "#327db7", "#6fafd2", "#c7e0ed",
                 "#fbd2bc", "#feab88", "#b71c2c", "#8b0823", "#6a0624"]
)

def auto_heatmap(
    data: pd.DataFrame,
    save_path: str,

    title: str = "",
    annot: bool = True,
    fmt: str = ".2f",
    figsize: tuple = (12, 10)
) -> None:
    """
    自动配色的相关系数热力图
    """
    # 根据数据规模自动选择配色
    n = data.shape[0]
    # cmap = custom_blue_purple
    # cmap = custom_orange_red
    # cmap = custom_green
    cmap = custom_blue_red
    draw_heatmap(
        data=data,
        title=title,
        save_path=save_path,
        cmap=cmap,
        annot=annot,
        fmt=fmt,
        figsize=figsize
    )



# 示例用法
if __name__ == "__main__":
    import numpy as np

    np.random.seed(42)
    types = ["高钾", "铅钡", "青铜器", "陶器"]
    components = ["SiO2", "CaO", "Fe2O3", "K2O"]
    test_data = pd.DataFrame({
        "成分": np.tile(components, 20),
        "变化百分比": np.random.normal(0, 20, 80),
        "文物类型": np.repeat(types, 20)
    })

    # auto_barplot(
    #     data=test_data,
    #     x_col="成分",
    #     y_col="变化百分比",
    #     hue_col="文物类型",
    #     title="自动配色-风化前后成分变化百分比",
    #     save_path="A/test_barplot_auto_palette.png"
    # )

    # auto_boxplot(
    #     data=test_data,
    #     x_col="成分",
    #     y_col="变化百分比",
    #     hue_col="文物类型",
    #     title="自动配色-成分变化分布箱线图",
    #     save_path="A/test_boxplot_auto_palette.png"
    # )

    # 随机生成相关矩阵
    np.random.seed(42)
    # mat = np.random.rand(8, 8)
    mat = np.random.uniform(-1, 1, (8, 8))  # 范围从-1到1
    df_heat = pd.DataFrame(mat, columns=[f"C{i+1}" for i in range(8)])

    # auto_heatmap(
    #     data=df_heat,
    #     title="自动配色-相关系数热力图",
    #     save_path="A/test_heatmap_auto_palette.png"
    # )

    # 测试饼状图
    # 创建文物类型分布数据
    artifact_types = ["高钾", "铅钡", "青铜器", "陶器", "其他"]
    counts = [25, 30, 20, 15, 10]
    pie_data = pd.Series(counts, index=artifact_types)

    # auto_pieplot(
    #     data=pie_data,
    #     title="文物类型分布",
    #     save_path="A/test_pieplot_auto_palette.png"
    # )

    # 测试带突出效果的饼状图
    auto_pieplot(
        data=pie_data[:4],  # 只取前4项
        title="主要文物类型分布",
        save_path="A/test_pieplot_explode.png",
        explode=[0.05, 0.1, 0, 0]  # 突出显示铅钡类型
    )